from langchain_core.documents.base import Document
from logging import Logger
import logging
from fastapi.applications import FastAPI
from typing import Union
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pathlib import Path
import os
from service.load_service import LoadService
from service.chunk_service import ChunkService


from fastapi import File, Form, UploadFile
from pydantic import BaseModel

#设置日志级别
logging.basicConfig(level=logging.INFO)
logger: Logger = logging.getLogger(__name__)

app: FastAPI = FastAPI(
    title="RAG",
    description="API文档",
    docs_url="/docs",  # Swagger UI路径 (默认是/docs)
)

# 确保必要的目录存在
os.makedirs("temp", exist_ok=True)
os.makedirs("loaded-docs", exist_ok=True)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 请求模型
class QueryRequest(BaseModel):
    question: str
    
    class Config:
        schema_extra = {
            "example": {
                "question": "什么是人工智能？",
            }
        }

# 响应模型
class QueryResponse(BaseModel):
    answer: str

    class Config:
        schema_extra = {
            "example": {
                "answer": "人工智能是...",
            }
        }

#请求
@app.post(path="/load",tags=["load"],description="加载文件接口")
async def process_file(
    file: UploadFile = File(default=None),
    web_url: str = Form(None),
    load_type: str = Form(),
    chunk_type: str = Form(default='character_text_splitter'),
    chunk_size: int = Form(default=100),
    chunk_overlap: int = Form(default=0)
    ):
    docs = []
    load_service: LoadService = LoadService()  
    #加载web网页
    if web_url and web_url.strip():
        docs = load_service.load_html(web_url=web_url,load_type=load_type)
    elif file:
        try:
            # 验证文件名
            if not file.filename:
                logger.error("文件名为空")
                return {"error": "文件名不能为空"}
            
            # 确保目录存在
            temp_dir = Path("temp")
            temp_dir.mkdir(exist_ok=True)
            
            # 保存上传的文件
            temp_path = temp_dir / file.filename
            with open(temp_path, "wb") as buffer:
                content = await file.read()
                buffer.write(content)
            
            logger.info(f"文件已保存到: {temp_path}")

            #根据文件类型加载
            file_extension = temp_path.suffix.lower()
            if file_extension == ".txt":
                docs = load_service.load_text(file_path=str(temp_path),load_type=load_type)
            elif file_extension == ".pdf":
                docs = load_service.load_pdf(file_path=str(temp_path),load_type=load_type)
            elif file_extension == ".csv":
                docs = load_service.load_csv(file_path=str(temp_path),load_type=load_type)
            elif file_extension == ".xlsx" or file_extension == ".xls":
                docs = load_service.load_excel(file_path=str(temp_path),load_type=load_type)
            elif file_extension == ".md":
                docs = load_service.load_markdown(file_path=str(temp_path),load_type=load_type)
            elif file_extension == ".jpg" or file_extension == ".png" or file_extension == ".jpeg":
                docs = load_service.load_image(file_path=str(temp_path),load_type=load_type)
        except Exception as e:
            logger.error(f"加载文件时出错: {str(e)}")
            return {"error": f"文件处理失败: {str(e)}"}
    
    if docs:
        chunk_service: ChunkService = ChunkService()
        chunks: list[Document] | None = chunk_service.chunk_docs(docs=docs,chunk_type=chunk_type,chunk_size=chunk_size,chunk_overlap=chunk_overlap)
        if chunks == None:
            return {
                "message": "切片后的数据为空"
            }
        elif chunks[0].metadata.get("error"):
            return {
                "message": chunks[0].page_content
            }
        else :
            return {
                "message": "切片处理成功", 
                "content": chunks,
                "load_type": load_type,
                "chunk_type": chunk_type,
                "chunk_size": chunk_size,
                "chunk_overlap": chunk_overlap
            }
    else:
        return {
            "message": "加载后的数据为空"
        }